package com.zwc.kotlin.coroutines.core

import com.zwc.kotlin.coroutines.Job
import com.zwc.kotlin.coroutines.OnCancel
import com.zwc.kotlin.coroutines.OnComplete
import com.zwc.kotlin.coroutines.scope.CoroutineScope
import java.util.concurrent.atomic.AtomicReference
import kotlin.coroutines.Continuation
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.resume
import kotlin.coroutines.suspendCoroutine

abstract class AbstractCoroutine<T>(context: CoroutineContext) : Job, Continuation<T>, CoroutineScope {
    protected val state = AtomicReference<CoroutineState>()
    override val context: CoroutineContext
    override val scopeContext: CoroutineContext
        get() = context

    init {
        state.set(CoroutineState.Incomplete())
        this.context = context + this
    }

    val isCompleted get() = state.get() is CoroutineState.Complete<*>
    override val isActive: Boolean
        get() = when (state.get()) {
            is CoroutineState.Cancelling,
            is CoroutineState.Complete<*> -> false
            else -> {
                true
            }
        }

    override suspend fun join() {
        when (state.get()) {
            is CoroutineState.Incomplete,
            is CoroutineState.Cancelling -> {
                return joinSuspend()
            }
            is CoroutineState.Complete<*> -> return
        }
    }

    private suspend fun joinSuspend() = suspendCoroutine<Unit> { continuation ->
        doOnCompleted { result ->
            continuation.resume(Unit)
        }
    }

    override fun resumeWith(result: Result<T>) {
        val newState = state.updateAndGet { prevState ->
            when (prevState) {
                is CoroutineState.Cancelling,
                is CoroutineState.Incomplete -> {
                    CoroutineState.Complete(result.getOrNull(), result.exceptionOrNull()).from(prevState)
                }
                is CoroutineState.Complete<*> -> {
                    throw IllegalStateException("Already completed!")
                }

            }
        }
        newState.notifyCompletion(result)
        newState.clear()
    }

    override fun invokeOnCancel(onCancel: OnCancel): Disposable {
        val disposable = CancellationHandlerDisposable(this, onCancel)
        val newState = state.updateAndGet { pre ->
            when (pre) {
                is CoroutineState.Incomplete -> {
                    CoroutineState.Incomplete().from(pre).with(disposable)
                }
                is CoroutineState.Cancelling,
                is CoroutineState.Complete<*> -> {
                    pre
                }
            }
        }
        (newState as? CoroutineState.Cancelling)?.let {
            onCancel()
        }
        return disposable
    }

    override fun invokeOnCompletion(onComplete: OnComplete): Disposable {
        return doOnCompleted {
            onComplete()
        }
    }

    override fun cancel() {
        val prevState = state.getAndUpdate { prev ->
            when (prev) {
                is CoroutineState.Incomplete -> CoroutineState.Cancelling()
                is CoroutineState.Cancelling,
                is CoroutineState.Complete<*> -> prev
            }
        }
        if (prevState is CoroutineState.Incomplete) {
            prevState.notifyCancellation()
            prevState.clear()
        }

    }

    override fun remove(disposable: Disposable) {
        state.updateAndGet { prev ->
            when (prev) {
                is CoroutineState.Incomplete -> {
                    CoroutineState.Incomplete().from(prev).without(disposable)
                }
                is CoroutineState.Cancelling -> {
                    CoroutineState.Cancelling().from(prev).without(disposable)
                }
                is CoroutineState.Complete<*> -> {
                    prev
                }
            }
        }
    }

    protected fun doOnCompleted(block: (Result<T>) -> Unit): Disposable {
        val disposable = CompletionHandlerDisposable(this, block)
        val newState = state.updateAndGet { prev ->
            when (prev) {
                is CoroutineState.Incomplete -> CoroutineState.Incomplete().from(prev).with(disposable)
                is CoroutineState.Cancelling -> CoroutineState.Cancelling().from(prev).with(disposable)
                is CoroutineState.Complete<*> -> prev
            }
        }
        (newState as? CoroutineState.Complete<T>)?.let {
            block(
                when {
                    it.exception != null -> Result.failure(it.exception)
                    it.value != null -> Result.success(it.value)
                    else -> {
                        throw IllegalStateException("Won't happen.")
                    }
                }
            )
        }
        return disposable
    }
}